# FileName:    main.py
# Date:        2022-09-22 15:27:40
# Author:      zzk
# Version:     1.0
# Description: re --> NFA


from typing import Dict, Set
import graphviz
from const import Const, StateType
from NFA import AdjMatrix
from typeassert import typeassert

adj_matrix = AdjMatrix()
conds = set()
nfa_ends = [] # record all end state

class Edge:
    def __init__(self, from_state, label, to_state=None) -> None:
        self.from_state = from_state # point to last State
        self.label = label
        self.to_state = to_state # point to next State

    def get_to_state_id(self):
        return self.to_state.id


class State:
    id_count = 0 # a static var
    def __init__(self) -> None:
        self.id = State.id_count # automatically increase state number
        State.id_count += 1
        self.out_edges = []      # one state could have a lot of out_edges, record all of them
        self.in_edges = []       # record every in_edge that point to self
        self.type = StateType.OTHER

# Three ways to check type
    # @typeassert(pnext=State)
    def add_edge(self, label, pnext):
        # if not isinstance(pnext, State): raise TypeError
        assert isinstance(pnext, State)
        edge = Edge(self, label, pnext)
        # maintain a matrix for nfa2dfa
        if adj_matrix[self.id].get(label) == None:
            adj_matrix[self.id][label] = []
        adj_matrix[self.id][label].append(pnext.id)
        if label != Const.EPSILON: conds.add(label) # record every cond except epsilon

        self.out_edges.append(edge)
        pnext.in_edges.append(edge) # record from edge

    def is_end(self):
        if self.type == StateType.END: return True
        return False

    def set_end(self):
        self.type = StateType.END
    
    def find_all_end(self):
        if self.is_end(): return
        self.set_end()
        nfa_ends.append(self.id)
        for inedge in self.in_edges:
            if inedge.label == Const.EPSILON:
                inedge.from_state.find_all_end() # TODO(zzk): Bad way to find end state
        for edge in self.out_edges:
            if edge.label == Const.EPSILON:
                edge.to_state.find_all_end()


class NFA:
    def __init__(self, start, end, label=Const.EPSILON) -> None:
        self.start = start
        self.end = end
        start.add_edge(label, self.end)

    def add_parallel_graph(self, nfa):
        newStart = State()
        newEnd = State()
        newStart.add_edge(Const.EPSILON, self.start)
        newStart.add_edge(Const.EPSILON, nfa.start)
        self.end.add_edge(Const.EPSILON, newEnd)
        nfa.end.add_edge(Const.EPSILON, newEnd)
        self.start = newStart
        self.end = newEnd

    def add_serial_graph(self, nfa):
        self.end.add_edge(Const.EPSILON, nfa.start)
        self.end = nfa.end

    def add_closure(self):
        self.end.add_edge(Const.EPSILON, self.start)
        self.start.add_edge(Const.EPSILON, self.end)

    def add_small_closure(self, ch):
        self.end.add_edge(ch, self.end)

    def add_one_node(self, label, pnext):
        # find edge repeat
        # for edge in self.end.edges:
        #     if edge.label == label and edge.ToState == pnext:
        #         return None
        self.end.add_edge(label, pnext)
        self.end = pnext
    
    def mark_all_end(self):
        self.end.find_all_end()

    def show(self, filename='nfa.gv'):
        self.mark_all_end()
        dirc = {}  # record every edge to avoid repeat paint
        stack = [] # Let's do some beautiful BFS
        g = graphviz.Digraph(filename=filename, engine='dot')
        g.attr('node', shape='circle')
        p = self.start
        while p is not None:
            for edge in p.out_edges:
                toStateID = edge.get_to_state_id()
                if toStateID not in dirc.get(p.id, []):
                    g.edge(str(p.id), str(toStateID), label=edge.label)
                    if edge.to_state.is_end(): g.body.append(f'{toStateID} [shape=doublecircle]') # Mark all end node(use double circle)
                    if dirc.get(p.id) == None:
                        dirc[p.id] = []
                    dirc[p.id].append(toStateID)
                    stack.append(edge.to_state)
            # The convenient of try block when BFS
            try:
                p = stack.pop()
            except IndexError:
                p = None
        g.view()


class DFAState:
    def __init__(self) -> None:
        self.id = 0 # the id is temp
        self.nfa_ids = []
        self.state_type = StateType.OTHER

    def __repr__(self) -> str:
        return str(self.id) + ':' + str(self.nfa_ids)

    def set(self, nfa_ids):
        self.nfa_ids = nfa_ids
    
    def get(self):
        return self.nfa_ids
    
    def extend(self, nfa_id):
        assert isinstance(nfa_id, list)
        for id in nfa_id:
            if id not in self.nfa_ids:
                self.nfa_ids.append(id)

    def compare(self, state_list):
        for idx, state in enumerate(state_list):
            if set(self.nfa_ids) == set(state.nfa_ids):
                return idx
        return -1
    
    def is_end(self):
        if self.state_type == StateType.END:
            return True
        return False

    def is_empty(self):
        if len(self.nfa_ids) == 0:
            return True
        return False    


class DFA:
    def __init__(self, condset) -> None:
        assert isinstance(condset, Set)
        self.states = []
        self.table = []
        self.condset = condset

    def __getitem__(self, id):
        assert isinstance(self.table[id], DFAState)
        return self.table[id]

    def add_state(self, dfa_state):
        assert isinstance(dfa_state, DFAState)
        # if the dfa state is end state
        for nfa_state in dfa_state.get():
            if nfa_state in nfa_ends:
                dfa_state.state_type = StateType.END
        self.states.append(dfa_state)

    def get_states(self):
        return self.states

    def get_conds(self):
        return self.condset

    def add_to_table(self, transform_dict):
        assert isinstance(transform_dict, Dict)
        self.table.append(transform_dict)

    def __repr__(self) -> str:
        return str(self.table)

    def show(self, filename='dfa.gv'):
        g = graphviz.Digraph(filename=filename, engine='dot')
        g.attr('node', shape='circle')
        for idx, transform_dict in enumerate(self.table):
            for key in transform_dict:
                g.edge(str(idx), str(transform_dict[key].id), label=str(key))
                if transform_dict[key].is_end():
                    g.body.append(f'{transform_dict[key].id} [shape=doublecircle]') # Mark all end node(use double circle)
        g.view()


# get substr between '(' ')', and return next index(ignore '(' and ')' )
# TODO(zzk): Do not support nested parentheses
def getsubregex(str, start):
    for i in range(start, len(str)):
        if str[i] == ')':
            return str[start + 1: i], i + 1
    assert 0 # illegal regex

# get substr between symbols, return index of next symbol
def getremainregex(str, start):
    for i in range(start + 1, len(str)):
        if str[i] in [')', '|', '*']:
            return str[start + 1: i], i
    return str[start + 1:], Const.MAX_RE_LEN

def checkclosure(nextch, nfa):
    # check if next char is *
    if nextch != '*':
        return False
    nfa.add_closure()
    return True

def re2nfa(str):
    substart = 0
    nfa = None
    while substart < len(str):
        ch = str[substart]
        if ch == '(': # get sub str
            substr, substart = getsubregex(str, substart)
            subnfa = re2nfa(substr)

            # TODO(zzk): Find a better way to deal with closure
            # bad code ==============================
            try: 
                if checkclosure(str[substart], subnfa): substart += 1 # if next char is *, skip it
            except IndexError: pass
            # bad code ==============================

            if nfa == None: nfa = subnfa
            else: nfa.add_serial_graph(subnfa) # be careful about this code
        elif ch == '|':
            substr, substart = getremainregex(str, substart)
            subnfa = re2nfa(substr)
            if nfa == None: nfa = subnfa
            else: nfa.add_parallel_graph(subnfa)
        else:
            if nfa == None: nfa = NFA(State(), State(), label=ch)
            else: 
                # bad code ==============================
                ifclosure = False
                try: 
                    if str[substart + 1] == '*':
                        ifclosure = True
                except IndexError: pass
                if ifclosure:
                    nfa.add_one_node(Const.EPSILON, State())
                    nfa.add_small_closure(ch)
                    substart += 1
                else:
                    nfa.add_one_node(ch, State())
                # bad code ==============================
            substart += 1
    # nfa.show()
    # sleep(2)
    return nfa

def getreachable(start_id, cond):
    new_ids = adj_matrix[start_id].get(cond, [])
    for id in new_ids:
        for newid in adj_matrix[id].get(cond, []):
            if newid not in new_ids: # remove duplicates
                new_ids.append(newid)
    # new_ids.append(start_id) # exclude start_id itself
    assert isinstance(new_ids, list)
    return new_ids

def nfa2dfa(nfa):
    # Subset Construction
    assert isinstance(nfa, NFA)
    # record all nfastates' epsilon closure!
    epsilon_closure = {}
    # find new start id in old nfa
    dfa_start_state = DFAState()
    epsilon_closure[nfa.start.id] = getreachable(nfa.start.id, Const.EPSILON)
    dfa_start_state.set(epsilon_closure[nfa.start.id] + [nfa.start.id]) # include start_id itself

    dfa = DFA(conds)
    dfa.add_state(dfa_start_state)
    for dfastate in dfa.get_states():
        transform_dict = {}
        for cond in dfa.get_conds():
            cond_to_state = DFAState()
            for nfa_state in dfastate.get():
                cond_to_state.extend(getreachable(nfa_state, cond))
                for nfa_cond_to_state in cond_to_state.get():
                    if epsilon_closure.get(nfa_cond_to_state) == None:
                        epsilon_closure[nfa_cond_to_state] = getreachable(nfa_cond_to_state, Const.EPSILON)
                    cond_to_state.extend(epsilon_closure[nfa_cond_to_state])
            
            # if cond_to_state is empty, do nothing
            if not cond_to_state.is_empty():
                # if the cond_to_state is a new dfa state, add it to dfa.states
                idx = cond_to_state.compare(dfa.get_states())
                if idx == -1:
                    cond_to_state.id = len(dfa.get_states())
                    dfa.add_state(cond_to_state)
                else:
                    cond_to_state = dfa.get_states()[idx]
                # store it to a temp dirc
                transform_dict[cond] = cond_to_state
        # deal with dfa transform table
        # print(transform_dict)
        # print(dfa)
        dfa.add_to_table(transform_dict) # even the dict is empty, it should be added to table
    return dfa

# You have to Mark All End State before using this function
# TODO(zzk): fatherid here is used to avoid infinite recursive loop. However, it can only avoid small circle(with just two State)
def is_match(string, pos, curState, fatherid):
    if pos >= len(string): 
        return curState.is_end() # recursive export
    for outedge in curState.out_edges:
        if outedge.label == Const.EPSILON and outedge.to_state.id != fatherid:
            if is_match(string, pos, outedge.to_state, curState.id): # try this way
                return True
        elif outedge.label == string[pos]:
            if is_match(string, pos + 1, outedge.to_state, curState.id):
                return True
    return False


# TODO(zzk): How to test my regex compiler efficiently?
testre = ['(a|b*)','(ab)*ca*', '(a|b|dd)c', '(a|b)*(aa|bb)(a|b)*']
teststr = ['a', 'abba', 'aaab']
# how to find end state?
# repeat edge!

def main():
    nfa = re2nfa(testre[0])
    nfa.show()
    # adj_matrix.show()
    dfa = nfa2dfa(nfa)
    print(nfa_ends)
    dfa.show()
    # print(is_match(teststr[2], 0, nfa.start, -1))


if __name__ == '__main__':
    main()
